"""
Author: yida
Time is: 2022/3/11 16:57 
this Code: 
"""
import os
import time
import torch
import torchvision.transforms as transforms
import torchvision.models.detection
import torchvision.models.detection.mask_rcnn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor, AnchorGenerator
# from torch.utils.tensorboard import SummaryWriter
 # 加载预训练模型
def get_squeezenet1_0_model_FRCNN(num_classes):
    '''
	backbone is efficientnet
	'''
    backbone = torchvision.models.squeezenet1_0(pretrained=True).features
    # backbone = EfficientNet.from_pretrained('efficientnet-b7')
    backbone.out_channels = 512
    anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),
                                aspect_ratios=((0.5, 1.0, 2.0),))
    # anchor_generator = AnchorGenerator(sizes=((8, 16, 32, 64, 128, 256, 512),),
	# 							aspect_ratios=((0.5, 1.0, 1,5, 2.0),))
    roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
													output_size=7,
													sampling_ratio=2)
    model = torchvision.models.detection.faster_rcnn.FasterRCNN(backbone,
					num_classes,
					rpn_anchor_generator=anchor_generator,
					box_roi_pool=roi_pooler)

    return  model
model = get_squeezenet1_0_model_FRCNN(12)
print(model)
with torch.no_grad(): # 在测试的时候必须加上这行和下一行代码，否则预测会出问题，这里是防止还有梯度更新这些，而且如果不加这个，后面又没有进行梯度更新的话，可能会报显存不够用的错误，我怀疑是数据没有被清理
    model.eval()
checkpoint = torch.load(r'/Users/yida/Desktop/31epoch/model_31.pth', map_location=torch.device('cpu'))
start_epoch=checkpoint['epoch']
model.load_state_dict(checkpoint['model'])
print(model)
#model.load_state_dict(torch.load(r'I:\test\coco\coco\checkpoints/model_31.pth'),False)#再加载网络的参数
model = model.to('cpu')
model.eval()
print(model)
inputs = torch.randn(1, 3, 800, 800)
outputs = model(inputs)
print(outputs)
# torch.onnx.export(model, (inputs), "./faster_rcnn.onnx", verbose=True)
# print("模型转换成功!")



inputs = torch. randn (1, 3, 800, 800, dtype=torch. float32, requires_grad=True).cpu()
inputs_2 = torch. randn (1, 3, 600, 800, dtype=torch. float32).cpu()
inputs =torch.tensor(inputs, dtype=torch. float32).cpu()
M = model (inputs)